# fastmcp_server.py
from __future__ import annotations
import json
import re
import os
import logging
import logging.handlers
from datetime import date, datetime
from functools import wraps
from typing import Optional, List, Literal, Any

from fastmcp import FastMCP
import mcp.types as types
from query_12306 import TrainInfoQuery

# ---------------- 日志配置 -----------------
LOG_LEVEL = os.getenv("LOG_LEVEL", "INFO").upper()
LOG_FMT = "%(asctime)s | %(levelname)8s | %(name)s | %(funcName)s:%(lineno)d | %(message)s"
os.makedirs("logs", exist_ok=True)
file_handler = logging.handlers.TimedRotatingFileHandler(
    filename="logs/mcp-12306-server.log",
    when="midnight",
    interval=1,
    backupCount=14,
    encoding="utf-8",
)
file_handler.setFormatter(logging.Formatter(LOG_FMT))
console_handler = logging.StreamHandler()
console_handler.setFormatter(logging.Formatter(LOG_FMT))
logging.basicConfig(level=LOG_LEVEL, handlers=[console_handler, file_handler], force=True)
logger = logging.getLogger("mcp-12306-server")

# ---------------- FastMCP 实例 -----------------
mcp = FastMCP("mcp-12306-server")

# ---------------- 工具异常回传装饰器 -----------------
def return_error(func: Any) -> Any:
    @wraps(func)
    async def wrapper(*args: Any, **kwargs: Any) -> List[types.ContentBlock]:
        try:
            return await func(*args, **kwargs)
        except Exception as exc:
            logger.exception("Tool %s failed", func.__name__)
            return [types.TextContent(type="text", text=f"[ERROR] {exc}")]
    return wrapper

# ---------------- 车站名→代码工厂 -----------------
def station_name_to_code_factory(query_engine: TrainInfoQuery) -> Any:
    def check_station(station: str) -> tuple[bool, Optional[str]]:
        if station in query_engine.station_name_dict:
            return True, station
        if station not in query_engine.city_station_dict:
            if station.endswith("站") or station.endswith("市"):
                station = station[:-1]
        if station in query_engine.city_station_dict or station in query_engine.station_name_dict:
            return True, station
        return False, None

    def convert(station_name: str) -> str:
        logger.debug("convert station_name=%s", station_name)
        name_is_valid, valid_station_name = check_station(station_name)
        if not name_is_valid or not valid_station_name:
            raise ValueError(f"[ERROR-STA-001] 无效车站名称: {station_name}，请检查是否输入简称或错别字")
        code = query_engine.get_station_code_by_name(valid_station_name)
        if not code:
            raise ValueError(f"[ERROR-STA-002] 未找到车站代码: {station_name}，请尝试去掉“站/市”后缀")
        logger.debug("station_name=%s -> code=%s", station_name, code)
        return code
    return convert

# ---------------- 复用引擎/转换器 -----------------
def _get_engine(func: Any) -> TrainInfoQuery:
    if not hasattr(func, "_engine"):
        func._engine = TrainInfoQuery()
    return func._engine

def _get_station_conv(func: Any) -> Any:
    if not hasattr(func, "_station_conv"):
        func._station_conv = station_name_to_code_factory(_get_engine(func))
    return func._station_conv

# ------------------------------------------------------
#  精准校验辅助函数（每种错误唯一标签）
# ------------------------------------------------------
def _validate_date(date_str: str) -> None:
    if not re.fullmatch(r"\d{4}-\d{2}-\d{2}", date_str):
        raise ValueError("[ERROR-DATE-001] 日期格式错误，应为 YYYY-MM-DD")
    try:
        dt = datetime.strptime(date_str, "%Y-%m-%d")
    except ValueError:
        raise ValueError("[ERROR-DATE-002] 无效日期，例：2025-01-01")
    if dt.date() < date.today():
        raise ValueError("[ERROR-DATE-003] 不能查询过去日期")
    if not TrainInfoQuery().is_date_can_query(date_str):
        raise ValueError("[ERROR-DATE-004] 只能查询未来 15 天内车票")

def _validate_purpose(codes: str) -> None:
    if codes not in {"ADULT", "0X00"}:
        raise ValueError("[ERROR-PUR-001] purpose_codes 只能为 ADULT 或 0X00（学生）")

def _validate_train_no(no: str) -> None:
    if not re.fullmatch(r"[GCDZTKYL]\d{1,4}", no, re.I):
        raise ValueError("[ERROR-NO-001] 车次号格式错误，例：G1、K123")

def _validate_station_code(code: str) -> None:
    if not re.fullmatch(r"[A-Z]{3}", code.upper()):
        raise ValueError("[ERROR-CODE-001] 车站代码必须为 3 位大写字母，例：VNP")

# =================================================================================
#  工具定义开始  ——  纯文本输出，无 Markdown
# =================================================================================

@mcp.tool()
@return_error
async def query_tickets(
    from_station: str,
    to_station: str,
    train_date: Optional[str] = None,
    purpose_codes: Literal["ADULT", "0X00"] = "ADULT",
) -> List[types.ContentBlock]:
    """
    查询 12306 指定日期、出发/到达站的余票信息（含各席别张数、车次、出发/到达时间）。

    参数
    ----
    from_station : str
        出发站中文名称，支持简称、模糊匹配（如“北京”≈“北京南”）。
    to_station : str
        到达站中文名称，规则同上。
    train_date : str, optional
        乘车日期，格式 YYYY-MM-DD，默认今天。
    purpose_codes : str, optional
        乘客类型，“ADULT”(成人) 或 “0X00”(学生)，默认“ADULT”。

    返回
    ----
    List[types.TextContent]
        纯文本，示例如下：
        2025-07-18  北京南 -> 上海虹桥  余票一览
        G1  06:00 -- 10:28  （04:28）
          商务 9张  一等 42张  二等 余票充足
    """
    engine = _get_engine(query_tickets)
    station_conv = _get_station_conv(query_tickets)
    date_str = train_date or date.today().strftime("%Y-%m-%d")
    _validate_date(date_str)
    _validate_purpose(purpose_codes)

    from_code = station_conv(from_station)
    to_code = station_conv(to_station)

    data = await engine.query_left_tickets(
        from_station=from_code,
        to_station=to_code,
        train_date=date_str,
        purpose_codes=purpose_codes,
    )
    if not data or not data.get("train_infos"):
        return [types.TextContent(type="text", text="该日期无车或已售完，请换日期或车站。")]

    lines = [f"{date_str}  {data['from_station']} -> {data['to_station']}  余票一览"]
    for t in data["train_infos"]:
        lines.append(f"{t['train_no']}  {t['start_time']} -- {t['arrive_time']}  （总时长{t['duration']}）")
        seat_parts = []
        for seat, info in t.get("seat_price", {}).items():
            seat_zh = {
                "business_seat": "商务",
                "premium_first_class_seat": "特等",
                "first_class_seat": "一等",
                "second_class_seat": "二等",
                "hard_sleep": "硬卧",
                "soft_sleep": "软卧",
                "hard_seat": "硬座",
                "no_seat": "无座",
            }.get(seat, seat)
            seat_parts.append(f"{seat_zh} {info['remain']}张")
        if seat_parts:
            lines.append("  " + "  ".join(seat_parts))
        else:
            lines.append("  暂无余票数据")
    return [types.TextContent(type="text", text="\n".join(lines))]


@mcp.tool()
@return_error
async def query_train_info(
    train_no: str,
    train_date: str,
) -> List[types.ContentBlock]:
    """
    查询指定车次在某天的完整时刻表（停靠站、到发时间、停车分钟）。

    参数
    ----
    train_no : str
        车次号，如“G1”、“K123”。
    train_date : str
        运行日期，格式 YYYY-MM-DD。

    返回
    ----
    List[types.TextContent]
        纯文本，示例如下：
        G1  2025-07-18 时刻表
        车站	到达	出发	停车
        北京南	--	06:00	--
        济南西	08:22	08:24	2
    """
    _validate_train_no(train_no)
    _validate_date(train_date)

    engine = _get_engine(query_train_info)
    data = await engine.query_train_info(train_no=train_no.upper(), train_date=train_date)
    if not data or not data.get("data"):
        return [types.TextContent(type="text", text="该车次指定日期未开行，请换日期或车次。")]

    stops = data["data"]
    lines = [f"{train_no.upper()}  {train_date} 时刻表"]
    lines.append("车站\t到达\t出发\t停车")
    for s in stops:
        lines.append(
            f"{s['station_name']}\t"
            f"{s.get('arrive_time') or '--'}\t"
            f"{s.get('start_time') or '--'}\t"
            f"{s.get('stopover_minutes') or '--'}"
        )
    return [types.TextContent(type="text", text="\n".join(lines))]


@mcp.tool()
@return_error
async def query_train_price(
    from_station: str,
    to_station: str,
    train_date: str,
    train_no: Optional[str] = None,
) -> List[types.ContentBlock]:
    """
    查询指定行程、日期的各席别票价（非余票）。

    参数
    ----
    from_station : str
        出发站中文名。
    to_station : str
        到达站中文名。
    train_date : str
        乘车日期，格式 YYYY-MM-DD。
    train_no : str, optional
        车次号，如"G1"、"K123"。如果不指定，返回所有车次的票价信息。

    返回
    ----
    List[types.TextContent]
        纯文本，示例如下：
        2025-07-18  北京南 -> 上海虹桥  各席别票价
        车次	席别	票价
        G1	二等座	¥553
    """
    _validate_date(train_date)
    engine = _get_engine(query_train_price)
    station_conv = _get_station_conv(query_train_price)

    from_code = station_conv(from_station)
    to_code = station_conv(to_station)

    data = await engine.query_train_price(
        train_no=train_no or "",               # 12306 接口实际无需 train_no，保留空串
        from_station=from_code,
        to_station=to_code,
        train_date=train_date,
        seat_types="O,1,2,3,4,5,6,7,8,9"  # 默认全部席别
    )
    if not data or not data.get("data"):
        return [types.TextContent(type="text", text="该行程无票价数据，请换车站或日期。")]

    price_info = data["data"]
    lines = [f"{train_date}  {from_station} -> {to_station}  各席别票价"]
    lines.append("车次\t席别\t票价")
    for seat_type, price in price_info.get("seat_types", {}).items():
        lines.append(f"{price_info['train_no']}\t{seat_type}\t¥{price}")
    return [types.TextContent(type="text", text="\n".join(lines))]


@mcp.tool()
@return_error
async def get_station_code(
    station_name: str,
) -> List[types.ContentBlock]:
    """
    根据车站中文名返回 12306 车站代码。

    参数
    ----
    station_name : str
        车站中文名，如“北京南”。

    返回
    ----
    List[types.TextContent]
        纯文本：北京南  的车站代码是：VNP
    """
    engine = _get_engine(get_station_code)
    code = engine.get_station_code_by_name(station_name)
    if not code:
        return [types.TextContent(type="text", text="未找到车站代码，请尝试去掉“站/市”后缀或检查错别字。")]
    return [types.TextContent(type="text", text=f"{station_name}  的车站代码是：{code}")]


@mcp.tool()
@return_error
async def get_station_name(
    station_code: str,
) -> List[types.ContentBlock]:
    """
    根据 12306 车站代码返回中文名。

    参数
    ----
    station_code : str
        车站代码，如"VNP"。

    返回
    ----
    List[types.TextContent]
        纯文本：车站代码 VNP 对应：北京南
    """
    _validate_station_code(station_code)
    engine = _get_engine(get_station_name)
    name = engine.get_station_name_by_code(station_code.upper())
    if not name:
        return [types.TextContent(type="text", text=f"[ERROR-CODE-002] 无效车站代码: {station_code.upper()}，请检查三位大写字母是否正确")]
    return [types.TextContent(type="text", text=f"车站代码 {station_code.upper()} 对应：{name}")]

# =================================================================================
#  启动入口
# =================================================================================
if __name__ == "__main__":
    logger.info("Starting mcp-12306-server on 0.0.0.0:7873")
    try:
        mcp.run(transport="streamable-http", host="0.0.0.0", port=7873)
    except TypeError:  # 旧版 fastmcp 回退
        import uvicorn
        uvicorn.run("fastmcp_server:mcp", host="0.0.0.0", port=7873)